import sys,os
import numpy as np
from math import pi, exp
from scipy.spatial.distance import euclidean,cdist

sigma=45

def kernelFunction(codeword, feature):
    # calcula kernel gaussiano usando a distancia euclidiana
    # K(D(w,ri)) = exp ( (1/2*pi*sigma^2) * D(ri,w)^2 )

    beta = -1/(2*pi*(sigma**2))
    dist = euclidean(codeword,feature)**2
    
    return exp(beta*dist)


# Coding step
def softCoding(features, codebook):

    # calcula o kernel gaussiano entre todas as features e codewords
    # gera uma matriz: codewords X features
    C = cdist(codebook, features, kernelFunction)

    # calcula o kernel codeword uncertainty para todas as features
    # UNC = K(D(w,ri))/ sum for all w [ K(D(w,ri) ]
    sumC = np.sum(C, axis=0)  # soma o kernel de uma feature para todas as codewords
    UNC = C/sumC

    return UNC


# Pooling Step
def averagePooling(C):
    return np.mean(C, axis=1)


# Pooling Step
def maxPooling(C):
    return np.amax(C, axis=1)


def readCodebook(filename):
# le o arquivo do codeboox e cria uma matriz de features com os descritores sift de cada codeword
    codebook = []

    fid = open(filename)
    lines = fid.readlines()
    fid.close()

    for line in lines:
        sl = line.split()
        codebook.append(sl[5:])

    return np.asarray(codebook)


# le a base, separando-a por imagem
def readBase(filename):
    base = dict()  # dicionario cujas chaves sao nomes de imagens da base

    fid = open(filename)
    lines = fid.readlines()
    fid.close()

    # cada imagem(chave) do dicionario 'base' recebera uma lista de features (sift dos pontos de interesse)
    for line in lines:
        sl = line.split()
        
        if base.has_key(sl[0]):
            base[sl[0]].append(np.array(sl[5:]))
        else:
            base[sl[0]] = [np.array(sl[5:])]
        
    return base


def main():

    # recebe o nome dos arquivos da base e do codebook
    baseFile = sys.argv[1]
    codebookFile = sys.argv[2] 

    # le codebook and base
    codebook = readCodebook(codebookFile)
    base = readBase(baseFile)

    # obtem o histograma descritor de cada imagem (bag of words)
    for k in base.keys():
        features = np.asarray(base[k])
        C = softCoding(features, codebook) # etapa de coding
        imgHist = maxPooling(C) # etapa de pooling

        # escreve o bag na saida padrao
        sys.stdout.write(k+ " ")
        sys.stdout.write(' '.join(str(x) for x in imgHist))
        sys.stdout.write('\n')


if __name__ == "__main__":
    main()
